package com.piece.core.db.util;

import com.piece.core.db.constant.SchemaConstants;
import org.springframework.core.io.InputStreamResource;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.EncodedResource;
import org.springframework.jdbc.datasource.DataSourceUtils;
import org.springframework.jdbc.datasource.init.ScriptException;
import org.springframework.jdbc.datasource.init.ScriptUtils;
import org.springframework.jdbc.datasource.init.UncategorizedScriptException;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import javax.sql.DataSource;
import java.io.IOException;
import java.io.InputStream;
import java.sql.Connection;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * 执行sql文件语句
 */
public class SchemaUtil {

    private SchemaUtil() {
        throw new IllegalStateException();
    }

    public static void execute(DataSource dataSource, List<InputStream> scripts) {
        if (CollectionUtils.isEmpty(scripts)) {
            return;
        }

        List<InputStream> locations = new ArrayList<>();
        scripts.forEach(script -> {
            locations.add(script);
        });

        initSchema(locations, dataSource);
        close(locations);
    }

    public static void initSchema(List<InputStream> locations, DataSource dataSource) {
        List<Resource> scripts = getResources(locations);
        if (!scripts.isEmpty()) {
            runScripts(dataSource, scripts);
        }
    }

    private static List<Resource> getResources(List<InputStream> locations) {
        List<Resource> resources = new ArrayList();
        Iterator iterator = locations.iterator();
        while (iterator.hasNext()) {
            InputStream inputStream = (InputStream) iterator.next();
            resources.add(getResource(inputStream));
        }
        return resources;
    }

    private static Resource getResource(InputStream inputStream) {
        try {
            return new InputStreamResource(inputStream);
        } catch (Exception ex) {
            throw new IllegalStateException("Unable to load resource", ex);
        }
    }

    private static void runScripts(DataSource dataSource, List<Resource> resources) {
        Assert.notNull(dataSource, "DataSource must not be null");
        try {
            Connection connection = DataSourceUtils.getConnection(dataSource);
            try {
                String encodedScript = SchemaConstants.sqlScriptEncoding;
                boolean continueOnError = SchemaConstants.continueOnError;
                String commentPrefix = "--";
                String separator = SchemaConstants.separator;
                String blockCommentStartDelimiter = "/*";
                String blockCommentEndDelimiter = "*/";
                resources.forEach(resource -> {
                    EncodedResource encodedResource = new EncodedResource(resource, encodedScript);
                    ScriptUtils.executeSqlScript(connection, encodedResource, continueOnError, true, commentPrefix, separator, blockCommentStartDelimiter, blockCommentEndDelimiter);
                });
            } finally {
                DataSourceUtils.releaseConnection(connection, dataSource);
            }
        } catch (Throwable e) {
            if (e instanceof ScriptException) {
                throw (ScriptException) e;
            } else {
                throw new UncategorizedScriptException("Failed to execute database script", e);
            }
        }
    }

    private static void close(List<InputStream> locations) {
        locations.stream().forEach(inputStream -> {
            if (null != inputStream) {
                try {
                    inputStream.close();
                } catch (IOException e) {
                }
            }
        });
    }
}
